package com.tuniu.agents.common.message;

import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import com.tuniu.agents.advisor.ChatEnvironmentMemory;
import com.tuniu.agents.advisor.ChatExtendMemoryAdvisor;
import com.tuniu.agents.agent.Agent;
import com.tuniu.agents.agent.AgentManager;
import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.common.pojo.ChatRecord;
import com.tuniu.agents.common.enums.MessageType;
import com.tuniu.agents.common.i18n.I18nConstants;
import com.tuniu.agents.common.message.adaper.AdapterText;
import com.tuniu.agents.common.service.AntispamService;
import com.tuniu.agents.common.service.ChatRecordService;
import com.tuniu.agents.util.I18nUtil;
import com.tuniu.agents.util.JsonUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

@Slf4j
@Component("conversationHistoryManager")
public class ConversationHistoryManager implements ConversationManager {

    @Autowired
    private ChatRecordService chatRecordService;

    @Autowired
    private AgentManager agentManager;

    @Autowired
    private ChatEnvironmentMemory chatEnvironmentMemory;

    @Autowired
    private ChatMemory chatMemory;

    /**
     * Antispam service
     */
    @Autowired(required = false)
    private AntispamService antispamService;

    /**
     * Conversation context for storing request and response content
     */
    private final Map<String, Map<String, List<String>>> CONVERSATION_CONTEXT = new ConcurrentHashMap<>();

    /**
     * Large model response plain text content key:responseId
     */
    private final Map<String,String> CONVERSATION_LLM_CONTEXT = new ConcurrentHashMap<>();

    /**
     * Set of message types that require content aggregation
     */
    private static final Set<String> AGGREGATABLE_TYPES = new HashSet<>(Arrays.asList("analysis", "text","footer"));


    @Override
    public void saveConversationRequest(String conversationId, String requestId, String responseId, String message, Integer isDeleted) {
        try {
            AdapterText text = new AdapterText().setContent(message)
                    .setType("text")
                    .setResponseId("");
            // Save user request to database
            String currentUserId = (String) chatEnvironmentMemory.get(conversationId, AgentsConstants.AgentsConversationEnvironmentKey.CUSTOMER_ID);
            if (StringUtils.isBlank(currentUserId)){
                currentUserId="";
            }
            chatRecordService.saveRequest(conversationId,currentUserId, requestId,responseId,JsonUtil.toJson(List.of(text)),isDeleted);
        } catch (Exception e) {
            log.error("Failed to save conversation request", e);
            throw new RuntimeException("Failed to save conversation request", e);
        }
    }

    @Override
    public String saveConversationResponse(String conversationId, String requestId, String responseId, String content) {
        if (responseId == null) {
            throw new IllegalStateException("responseId is not set for the current request");
        }
        // Parse response content
        if (!JsonUtil.isValidJson(content)) {
            appendToLlmContext(responseId,content);
            // Invalid JSON, insert into main response information
            content = BaseAgentNewResp.sendAnswerMsg(responseId, content);
        }else{
            appendToLlmContext(responseId,JsonUtil.fromJson(content,BaseAgentNewResp.class).getContent());
        }
        appendToStream(conversationId, responseId, content);
        return content;
    }

    @Override
    public String saveConversationFooter(String conversationId, String requestId, String responseId, String content) {
        if (responseId == null) {
            throw new IllegalStateException("responseId is not set for the current request");
        }
        content = BaseAgentNewResp.sendFooterMsg(responseId, initFooterMsgWithModel(conversationId));
        appendToLlmContext(responseId,content);
        appendToStream(conversationId, responseId, content);
        return content;
    }

    private String initFooterMsgWithModel(String conversationId){
        String msg = I18nUtil.getMessage(conversationId,I18nConstants.FOOTER_CONTENT_MESSAGE_KEY);
        if(StringUtils.isNotBlank(conversationId)){
            Object modelObject = chatEnvironmentMemory.get(conversationId, ChatExtendMemoryAdvisor.CHAT_MODEL_LIST);
            if(modelObject!=null) {
                List<String> modelName = JsonUtil.toList(modelObject.toString(), String.class);
                // Convert modelName to a string, separated by commas
                String modelNameStr = ": "+String.join(",", modelName);
                msg = I18nUtil.getMessage(conversationId,I18nConstants.FOOTER_CONTENT_MESSAGE_KEY, new Object[]{modelNameStr});
            }
        }
        chatEnvironmentMemory.clear(conversationId,ChatExtendMemoryAdvisor.CHAT_MODEL_LIST);
        return msg;
    }

    @Override
    public String saveConversationResponse(String conversationId, String content) {
        List<ChatRecord> chatRecords = chatRecordService.findByConversationIdAndMessageTypeOrderByAddTimeDesc(conversationId, MessageType.REQUEST.getCode());
        ChatRecord chatRecord = chatRecords.get(0);
        try {
            String currentUserId = (String) chatEnvironmentMemory.get(conversationId, AgentsConstants.AgentsConversationEnvironmentKey.CUSTOMER_ID);
            if (StringUtils.isBlank(currentUserId)){
                currentUserId="";
            }
            // Save response to database
            chatRecordService.savePrimaryResponse(currentUserId,conversationId,chatRecord.getRequestId(), chatRecord.getResponseId(), content,0);
        } catch (Exception e) {
            log.error("Failed to save conversation response", e);
            throw new RuntimeException("Failed to save conversation response", e);
        }
        return null;
    }

    private void appendToStream(String conversationId, String responseId, String content) {
        // Get or create LinkedHashMap corresponding to conversationId
        Map<String, List<String>> responseMap = CONVERSATION_CONTEXT.computeIfAbsent(conversationId, k -> new LinkedHashMap<>());
        List<String> respList = responseMap.computeIfAbsent(responseId, k -> new CopyOnWriteArrayList<>());
        // Serialize content into encapsulation class
        JsonObject resp = JsonUtil.toJsonObject(content);
        String cResponseId = getResponseIdFromJson(resp);
        if (StringUtils.isBlank(cResponseId)){
            resp.addProperty("responseId", responseId);
            respList.add(JsonUtil.toJson(resp));
        }else {
            List<String> childRespList = responseMap.computeIfAbsent(cResponseId, k -> new CopyOnWriteArrayList<>());
            childRespList.add(JsonUtil.toJson(resp));
        }
    }

    @Override
    public Flux<String> getQuestions(String conversationId) {
        if (StringUtils.isNotEmpty(conversationId)) {
            // Retrieve conversation history from the database
            List<ChatRecord> history = chatRecordService.getConversationHistory(conversationId);
            Map<String, Object> conversationMap = new HashMap<>();

            // Convert to original format
            for (ChatRecord record : history) {
                Map<String, Object> content = chatRecordService.parseContent(record.getContent());
                if (record.getMessageType().equals(MessageType.REQUEST.getCode())) {
                    conversationMap.put(record.getRequestId(), content.get("message"));
                } else if (record.getMessageType().equals(MessageType.PRIMARY_RESPONSE.getCode())) {
                    conversationMap.put(record.getResponseId(), content.get("answer"));
                }
            }

            // Use corresponding agent
            Agent agent = agentManager.getAgent("QUESTIONS");
            return agent.stream("", JsonUtil.toJson(conversationMap));
        }
        return Flux.empty();
    }

    /**
     * Save streaming response content to the database
     */
    @Override
    public void streamToDatabase(String conversationId, String requestId, String responseId) {
        try {
            if (requestId == null) {
                log.error("Cannot find request ID for conversation: {}", conversationId);
                return;
            }

            Map<String, List<String>> responseMap = CONVERSATION_CONTEXT.get(conversationId);
            if (MapUtils.isEmpty(responseMap)) {
                log.error("Cannot find response content for conversation: {}", conversationId);
                return;
            }
            Map<String, JsonArray> aggregatedResponseMap = new HashMap<>();

            // Aggregate operation for each responseId
            for (Map.Entry<String, List<String>> entry : responseMap.entrySet()) {
                String currentResponseId = entry.getKey();
                List<String> respList = entry.getValue();

                // Aggregate operation, preserving original order
                JsonArray aggregatedList = aggregateResponsesPreservingOrder(respList);
                aggregatedResponseMap.put(currentResponseId, aggregatedList);
            }
            // Save aggregated responses
            for (Map.Entry<String, JsonArray> entry : aggregatedResponseMap.entrySet()) {
                if (entry.getKey().equals(responseId)) {
                    String currentUserId = (String) chatEnvironmentMemory.get(conversationId, AgentsConstants.AgentsConversationEnvironmentKey.CUSTOMER_ID);
                    if (StringUtils.isBlank(currentUserId)){
                        currentUserId="";
                    }
                    // Determine if logical storage is needed
                    if( chatEnvironmentMemory.get(conversationId,entry.getKey())!=null && chatEnvironmentMemory.get(conversationId,entry.getKey()).toString().equals("1")) {
                        chatRecordService.savePrimaryResponse(currentUserId, conversationId, requestId, entry.getKey(), JsonUtil.toJson(entry.getValue()),1);
                        chatEnvironmentMemory.clear(conversationId);
                    }else {
                        chatRecordService.savePrimaryResponse(currentUserId, conversationId, requestId, entry.getKey(), JsonUtil.toJson(entry.getValue()),0);
                    }
                } else {
                    // Determine if logical storage is needed
                    if( chatEnvironmentMemory.get(conversationId,responseId)!=null && chatEnvironmentMemory.get(conversationId,responseId).toString().equals("1")) {
                        chatRecordService.saveSecondaryResponse(requestId, entry.getKey(), responseId, JsonUtil.toJson(entry.getValue()),1);
                        chatEnvironmentMemory.clear(conversationId);
                    }else {
                        chatRecordService.saveSecondaryResponse(requestId, entry.getKey(), responseId, JsonUtil.toJson(entry.getValue()),0);
                    }
                }
            }

        } catch (Exception e) {
            log.error("Failed to save streaming response to database", e);
        } finally {
            // Clear temporary storage
            CONVERSATION_CONTEXT.remove(conversationId);
        }
    }

     public String checkAntispam(String conversationId, String requestId, String responseId, String content){
        if(antispamService == null) {
            return content;
        }
        try {
            String checkText = CONVERSATION_LLM_CONTEXT.get(responseId);
            if (StringUtils.isEmpty(antispamService.antispamCheck(conversationId, checkText, responseId))) {
                // Check passed
                return BaseAgentNewResp.sendAnswerMsg(responseId,"");
            } else {
                chatEnvironmentMemory.add(conversationId,Map.of(responseId,1));
                chatMemory.clear(conversationId);
                // Check failed
                return content;
            }
        }catch (Exception e){
            log.error("Validation failed", e);
            return content;
        }finally {
            CONVERSATION_LLM_CONTEXT.remove(responseId);
        }
    }

    /**
     * Aggregate responses based on the type field, preserving the original order, considering different content structures
     * @param respList Original response list
     * @return Aggregated response list
     */
    private JsonArray aggregateResponsesPreservingOrder(List<String> respList) {
        if (respList == null || respList.isEmpty()) {
            return new JsonArray();
        }

        JsonArray result = new JsonArray();
        int i = 0;

        while (i < respList.size()) {

            JsonObject current = JsonUtil.toJsonObject(respList.get(i));
            String currentType = getTypeFromJson(current);
            // Check if the current type needs aggregation
            if (AGGREGATABLE_TYPES.contains(currentType)) {
                StringBuilder combinedContent = new StringBuilder(getContentFromJson(current));
                int j = i + 1;
                // Find consecutive messages of the same type for aggregation
                while (j < respList.size() &&
                        currentType.equals(getTypeFromJson(JsonUtil.toJsonObject(respList.get(j)))) &&
                        AGGREGATABLE_TYPES.contains(getTypeFromJson(JsonUtil.toJsonObject(respList.get(j))))) {
                    combinedContent.append(getContentFromJson(JsonUtil.toJsonObject(respList.get(j))));
                    j++;
                }

                // Create aggregated message
                if (j > i + 1) {  // Aggregation occurred
                    current.addProperty("content", combinedContent.toString());
                    result.add(current);
                    i = j;  // Skip aggregated messages
                } else {
                    // No aggregation possible, add current message directly
                    result.add(current);
                    i++;
                }
            } else {
                // Type does not require aggregation, add directly
                result.add(current);
                i++;
            }
        }
        return result;

    }

    public void appendToLlmContext(String key, Object text) {
        if(text == null){
            return;
        }
        // Use the compute method to handle cases where the key exists or does not exist
        CONVERSATION_LLM_CONTEXT.compute(key, (k, oldValue) -> {
            if (oldValue == null) {
                // If the key does not exist, initialize as new text
                return text.toString();
            } else {
                // If the key exists, append text to the existing value
                return oldValue + text.toString();
            }
        });
    }

    private String getFieldFromJson(JsonObject jsonObject,String key) {
        JsonElement jsonElement = JsonUtil.getJsonValue(jsonObject, key);
        if (jsonElement == null){
            return "";
        }
        return jsonElement.getAsString();
    }

    private String getContentFromJson(JsonObject jsonObject) {
        return getFieldFromJson(jsonObject,"content");
    }

    private String getTypeFromJson(JsonObject jsonObject) {
        return getFieldFromJson(jsonObject,"type");
    }

    private String getResponseIdFromJson(JsonObject jsonObject) {
        return getFieldFromJson(jsonObject,"responseId");
    }
}